import dgl
import torch


# 实现采样和子图的提取
# 注意加入了总图graph_all的处理
class MinibatchSampler(object):
    def __init__(self, graph_score, graph_commit, graph_collect, graph_duration, graph_all, num_layers):
        self.graph_score = graph_score
        self.graph_commit = graph_commit
        self.graph_collect = graph_collect
        self.graph_duration = graph_duration
        self.graph_all = graph_all
        self.num_layers = num_layers

    def sample(self, batch):
        ps = []
        bs = []

        users, items, rating_score, rating_commit, rating_collect, rating_duration, rating_all = zip(*batch)
        users = torch.stack(users)
        items = torch.stack(items)
        rating_score = torch.stack(rating_score)
        rating_commit = torch.stack(rating_commit)
        rating_collect = torch.stack(rating_collect)
        rating_duration = torch.stack(rating_duration)
        rating_all = torch.stack(rating_all)
        # --------------------------------------------------score
        # 1 创建二部图
        pair_graph_score = dgl.heterograph(
            {('user', 'watched', 'item'): (users, items)},
            num_nodes_dict={'user': self.graph_score.num_nodes('user'),
                            'item': self.graph_score.num_nodes('item')}
        )
        u = users.tolist()
        i = items.tolist()
        real_data = torch.tensor(list(zip(u, i)), dtype=torch.int)
        pair_graph_score.edata['real_data'] = real_data

        # 2 压缩二部图
        pair_graph_score = dgl.compact_graphs(pair_graph_score)
        pair_graph_score.edata['rating'] = rating_score

        # 3 创建数据块
        seeds = {'user': pair_graph_score.nodes['user'].data[dgl.NID],
                 'item': pair_graph_score.nodes['item'].data[dgl.NID]}
        block_score = self.construct_blocks(seeds, (users, items), self.graph_score)

        # 把节点特征也复制过来
        # 注意这里只需要处理源端结点
        for feature_name in self.graph_score.nodes['user'].data.keys():
            block_score[0].srcnodes['user'].data[feature_name] = \
                self.graph_score.nodes['user'].data[feature_name][block_score[0].srcnodes['user'].data[dgl.NID]]
        for feature_name in self.graph_score.nodes['item'].data.keys():
            block_score[0].srcnodes['item'].data[feature_name] = \
                self.graph_score.nodes['item'].data[feature_name][block_score[0].srcnodes['item'].data[dgl.NID]]
        ps.append(pair_graph_score)
        bs.append(block_score)
        # -------------------------------------------------commit
        # 1 创建二部图
        pair_graph_commit = dgl.heterograph(
            {('user', 'watched', 'item'): (users, items)},
            num_nodes_dict={'user': self.graph_commit.num_nodes('user'),
                            'item': self.graph_commit.num_nodes('item')}
        )
        u = users.tolist()
        i = items.tolist()
        real_data = torch.tensor(list(zip(u, i)), dtype=torch.int)
        pair_graph_commit.edata['real_data'] = real_data

        # 2 压缩二部图
        pair_graph_commit = dgl.compact_graphs(pair_graph_commit)
        pair_graph_commit.edata['rating'] = rating_commit

        # 3 创建数据块
        seeds = {'user': pair_graph_commit.nodes['user'].data[dgl.NID],
                 'item': pair_graph_commit.nodes['item'].data[dgl.NID]}
        block_commit = self.construct_blocks(seeds, (users, items), self.graph_commit)

        # 把节点特征也复制过来
        # 注意这里只需要处理源端结点
        for feature_name in self.graph_commit.nodes['user'].data.keys():
            block_commit[0].srcnodes['user'].data[feature_name] = \
                self.graph_commit.nodes['user'].data[feature_name][block_commit[0].srcnodes['user'].data[dgl.NID]]
        for feature_name in self.graph_commit.nodes['item'].data.keys():
            block_commit[0].srcnodes['item'].data[feature_name] = \
                self.graph_commit.nodes['item'].data[feature_name][block_commit[0].srcnodes['item'].data[dgl.NID]]
        ps.append(pair_graph_commit)
        bs.append(block_commit)
        # -------------------------------------------------collect
        # 1 创建二部图
        pair_graph_collect = dgl.heterograph(
            {('user', 'watched', 'item'): (users, items)},
            num_nodes_dict={'user': self.graph_collect.num_nodes('user'),
                            'item': self.graph_collect.num_nodes('item')}
        )
        u = users.tolist()
        i = items.tolist()
        real_data = torch.tensor(list(zip(u, i)), dtype=torch.int)
        pair_graph_collect.edata['real_data'] = real_data

        # 2 压缩二部图
        pair_graph_collect = dgl.compact_graphs(pair_graph_collect)
        pair_graph_collect.edata['rating'] = rating_collect

        # 3 创建数据块
        seeds = {'user': pair_graph_collect.nodes['user'].data[dgl.NID],
                 'item': pair_graph_collect.nodes['item'].data[dgl.NID]}
        block_collect = self.construct_blocks(seeds, (users, items), self.graph_collect)

        # 把节点特征也复制过来
        # 注意这里只需要处理源端结点
        for feature_name in self.graph_collect.nodes['user'].data.keys():
            block_collect[0].srcnodes['user'].data[feature_name] = \
                self.graph_collect.nodes['user'].data[feature_name][block_collect[0].srcnodes['user'].data[dgl.NID]]
        for feature_name in self.graph_collect.nodes['item'].data.keys():
            block_collect[0].srcnodes['item'].data[feature_name] = \
                self.graph_collect.nodes['item'].data[feature_name][block_collect[0].srcnodes['item'].data[dgl.NID]]
        ps.append(pair_graph_collect)
        bs.append(block_collect)
        # -------------------------------------------------duration
        # 1 创建二部图
        pair_graph_duration = dgl.heterograph(
            {('user', 'watched', 'item'): (users, items)},
            num_nodes_dict={'user': self.graph_duration.num_nodes('user'),
                            'item': self.graph_duration.num_nodes('item')}
        )
        u = users.tolist()
        i = items.tolist()
        real_data = torch.tensor(list(zip(u, i)), dtype=torch.int)
        pair_graph_duration.edata['real_data'] = real_data

        # 2 压缩二部图
        pair_graph_duration = dgl.compact_graphs(pair_graph_duration)
        pair_graph_duration.edata['rating'] = rating_duration

        # 3 创建数据块
        seeds = {'user': pair_graph_duration.nodes['user'].data[dgl.NID],
                 'item': pair_graph_duration.nodes['item'].data[dgl.NID]}
        block_duration = self.construct_blocks(seeds, (users, items), self.graph_duration)

        # 把节点特征也复制过来
        # 注意这里只需要处理源端结点
        for feature_name in self.graph_duration.nodes['user'].data.keys():
            block_duration[0].srcnodes['user'].data[feature_name] = \
                self.graph_duration.nodes['user'].data[feature_name][block_duration[0].srcnodes['user'].data[dgl.NID]]
        for feature_name in self.graph_duration.nodes['item'].data.keys():
            block_duration[0].srcnodes['item'].data[feature_name] = \
                self.graph_duration.nodes['item'].data[feature_name][block_duration[0].srcnodes['item'].data[dgl.NID]]
        ps.append(pair_graph_duration)
        bs.append(block_duration)

        # -----------------总图all的处理----------------------
        # 1 创建二部图
        pair_graph_all = dgl.heterograph(
            {('user', 'watched', 'item'): (users, items)},
            num_nodes_dict={'user': self.graph_all.num_nodes('user'),
                            'item': self.graph_all.num_nodes('item')}
        )
        u = users.tolist()
        i = items.tolist()
        real_data = torch.tensor(list(zip(u, i)), dtype=torch.int)
        pair_graph_all.edata['real_data'] = real_data

        # 2 压缩二部图
        pair_graph_all = dgl.compact_graphs(pair_graph_all)
        pair_graph_all.edata['rating'] = rating_all

        # 3 创建数据块
        seeds = {'user': pair_graph_all.nodes['user'].data[dgl.NID],
                 'item': pair_graph_all.nodes['item'].data[dgl.NID]}
        block_all = self.construct_blocks(seeds, (users, items), self.graph_all)

        # 把节点特征也复制过来
        # 注意这里只需要处理源端结点
        for feature_name in self.graph_all.nodes['user'].data.keys():
            block_all[0].srcnodes['user'].data[feature_name] = \
                self.graph_all.nodes['user'].data[feature_name][block_all[0].srcnodes['user'].data[dgl.NID]]
        for feature_name in self.graph_all.nodes['item'].data.keys():
            block_all[0].srcnodes['item'].data[feature_name] = \
                self.graph_all.nodes['item'].data[feature_name][block_all[0].srcnodes['item'].data[dgl.NID]]
        ps.append(pair_graph_all)
        bs.append(block_all)

        return ps, bs

    def construct_blocks(self, seeds, user_item_pairs_to_remove, graph):
        blocks = []
        users, items = user_item_pairs_to_remove
        # 采样就是根据卷积层数选取对应数量的邻居结点
        # 涉及到双向图的处理
        for i in range(self.num_layers):
            sampled_graph = dgl.in_subgraph(graph, seeds)
            sampled_eids = sampled_graph.edges[('user', 'watched', 'item')].data[dgl.EID]
            sampled_eids_rev = sampled_graph.edges[('item', 'watchedby', 'user')].data[dgl.EID]

            # 训练时要去掉用户和项目间的关联
            _, _, edges_to_remove = sampled_graph.edge_ids(
                users, items, etype=('user', 'watched', 'item'), return_uv=True)
            _, _, edges_to_remove_rev = sampled_graph.edge_ids(
                items, users, etype=('item', 'watchedby', 'user'), return_uv=True)

            # sampled_with_edges_removed = dgl.remove_edges(
            #     sampled_graph,
            #     {('user', 'watched', 'item'): edges_to_remove, ('item', 'watchedby', 'user'): edges_to_remove_rev}
            # )

            sampled_with_edges_removed = dgl.remove_edges(sampled_graph,
                                                          edges_to_remove, ('user', 'watched', 'item'))
            sampled_with_edges_removed = dgl.remove_edges(sampled_with_edges_removed,
                                                          edges_to_remove_rev, ('item', 'watchedby', 'user'))

            sampled_eids = sampled_eids[
                sampled_with_edges_removed.edges[('user', 'watched', 'item')].data[dgl.EID]]
            sampled_eids_rev = sampled_eids_rev[
                sampled_with_edges_removed.edges[('item', 'watchedby', 'user')].data[dgl.EID]]

            # 创建子图块
            block = dgl.to_block(sampled_with_edges_removed, seeds)
            blocks.insert(0, block)
            seeds = {'user': block.srcnodes['user'].data[dgl.NID],
                     'item': block.srcnodes['item'].data[dgl.NID]
                     }

            # 把评分复制过去
            block.edges[('user', 'watched', 'item')].data['rating'] = \
                graph.edges[('user', 'watched', 'item')].data['rating'][sampled_eids]
            block.edges[('item', 'watchedby', 'user')].data['rating'] = \
                graph.edges[('item', 'watchedby', 'user')].data['rating'][sampled_eids_rev]

        return blocks
